import asyncio
from datetime import datetime
from beanie import Document, init_beanie
from motor.motor_asyncio import AsyncIOMotorClient
from beanie.operators import CurrentDate, Set
from motor.motor_asyncio import AsyncIOMotorCollection
from pymongo import InsertOne, UpdateOne, DeleteOne
from conn import mongo_url

client = AsyncIOMotorClient(mongo_url)
db = client["test"]


# 定义文档模型
class User(Document):
    name: str
    updated_at: datetime | None = None


class Book(Document):
    title: str
    author: str


async def test_current_date():
    async with await client.start_session() as session, session.start_transaction():
        cnt = await User.find(session=session).count()
        # 创建一个新用户
        user = User(name="John Doe")
        await user.insert(session=session)

        # 使用 $currentDate 更新 last_login 字段
        await User.find_one(User.id == user.id).update(
            CurrentDate({User.updated_at: True}),
            session=session,
        )

        # 重新获取用户文档以查看更新后的值
        updated_user = await User.get(user.id, session=session)
        print(f"User: {updated_user.name}, Updated At: {updated_user.updated_at}")

        await User.find(User.id == updated_user.id).delete()

        new_cnt = await User.find(session=session).count()
        assert cnt == new_cnt, f"{cnt} != {new_cnt}"


async def test_bulk_write():
    collection: AsyncIOMotorCollection = Book.get_motor_collection()

    # 定义批量写入的操作

    async with await client.start_session() as session, session.start_transaction():
        operations = [
            InsertOne({"title": "Book 1", "author": "Author 1"}),
            InsertOne({"title": "Book 2", "author": "Author 2"}),
            # UpdateOne({"title": "Book 1"}, Set({"author": "New Author"})),
            UpdateOne({Book.title: "Book 1"}, Set({Book.author: "New Author"})),
            DeleteOne({"title": "Book 2"}),
        ]

        # 执行批量写入操作
        result = await collection.bulk_write(operations, session=session)

        print(f"Inserted: {result.inserted_count}")
        print(f"Modified: {result.modified_count}")
        print(f"Deleted: {result.deleted_count}")


async def main():
    # 初始化 Beanie
    await init_beanie(database=db, document_models=[User, Book])

    # await test_current_date()
    await test_bulk_write()


if __name__ == "__main__":
    asyncio.run(main())
